#include <stdio.h>
#include <stdlib.h>
#include <openssl/sha.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#define MALLOC_SIGN_SIZE 256
#define BOOT_IMG_SIZE 10485760

// Utility function to read the entire file into memory
unsigned char* read_file(const char* filename, size_t* size) {
    FILE* file = fopen(filename, "rb");
    if (!file) {
        perror("Error opening file");
        return NULL;
    }

    fseek(file, 0, SEEK_END);
    *size = ftell(file);
    fseek(file, 0, SEEK_SET);

    unsigned char* buffer = (unsigned char*)malloc(*size);
    if (!buffer) {
        perror("Memory allocation failed");
        fclose(file);
        return NULL;
    }

    fread(buffer, 1, *size, file);
    fclose(file);
    return buffer;
}

// The main function which verifies the signature
int main() {
    size_t img_size, sign_size, key_size;
    unsigned char* boot_img = read_file("boot.img", &img_size);
    unsigned char* boot_sign = read_file("boot.signature", &sign_size);
    unsigned char* public_key_data = read_file("public.pem", &key_size);
    img_size -= sign_size;
    if (!boot_img || !boot_sign || !public_key_data) {
        // Handle error...
        free(boot_img);
        free(boot_sign);
        free(public_key_data);
        return 1;
    }

    // Extract the public key
    RSA* public_key = RSA_new();
    BIO* bio = BIO_new_mem_buf(public_key_data, key_size);
    PEM_read_bio_RSA_PUBKEY(bio, &public_key, NULL, NULL);
    BIO_free(bio);
    free(public_key_data);

    // Create SHA256 hash of the image
    unsigned char hash[SHA256_DIGEST_LENGTH];
    SHA256(boot_img, img_size, hash);
    free(boot_img);

    unsigned char *final_sign = (unsigned char *) malloc(MALLOC_SIGN_SIZE);
    if (final_sign == NULL) {
        printf("malloc final_sign error\n");
        return 1;
    }

    FILE *fp;
    fp = fopen("boot.img","rb");
    fseek(fp,BOOT_IMG_SIZE,SEEK_SET);
    fread(final_sign, MALLOC_SIGN_SIZE, 1, fp);
    fclose(fp);
    // Verify the signature
    int result = RSA_verify(NID_sha256, hash, SHA256_DIGEST_LENGTH, final_sign, MALLOC_SIGN_SIZE, public_key);
    free(final_sign);
    free(boot_sign);
    RSA_free(public_key);
    if (result != 1) {
        fprintf(stderr, "Signature verification failed: %s\n", ERR_error_string(ERR_get_error(), NULL));
        return 2;
    }
    printf("Signature is valid.\n");
    return 0;
}